from PyQt5.QtCore import QThread

from model.dataloaders import LoadSource
from model.detect import SingleDetect

Window = None


class DetectThread(QThread):
    def __init__(self, source, parent: Window = None, img_size=640, stride=32, auto=True):
        super().__init__(parent)
        self.parent = parent
        self.source = source
        self.dataloader = LoadSource(source, img_size=img_size, stride=stride, auto=auto)

    def is_running(self):
        return self.dataloader.running

    def stop(self):
        self.dataloader.stop()


class VideoThread(DetectThread):
    def __init__(self, source, parent: Window = None, img_size=640, stride=32, auto=True):
        super().__init__(source, parent, img_size=img_size, stride=stride, auto=auto)

    def run(self):
        opened = False
        for img in self.dataloader:
            # img, img0 = item[0], item[1]
            ret = self.parent.detector(img)
            if opened is False:
                opened = True
                self.parent.stop_loading_signal.emit('视频已打开')
            self.parent.show_detect_result.emit(ret)
        self.parent.show_detect_result.emit("")


class ImageThread(DetectThread):
    def __init__(self, source, parent: Window = None, img_size=800, stride=32, auto=True):
        super().__init__(source, parent, img_size=img_size, stride=stride, auto=auto)

    def run(self):
        for img in self.dataloader:
            # img, img0 = item[0], item[1]
            ret = self.parent.detector(img)
            self.parent.stop_loading_signal.emit('检测完成')
            self.parent.show_detect_result.emit(ret)


class CammeraThread(DetectThread):
    def __init__(self, source, parent: Window = None, img_size=640, stride=32, auto=True):
        super().__init__(source, parent, img_size=img_size, stride=stride, auto=auto)

    def run(self):
        opened = False
        for img in self.dataloader:
            # img, img0 = item[0], item[1]
            ret = self.parent.detector(img)
            if opened is False:
                opened = True
                self.parent.stop_loading_signal.emit('摄像头已打开')
            self.parent.show_detect_result.emit(ret)
        self.parent.show_detect_result.emit("")


class WarmupThread(QThread):

    def __init__(self, parent: Window,
                 path, device, dtype,
                 iou_thres, conf_thres,
                 imgsz=640):
        super().__init__(parent=parent)
        self.parent = parent
        self.path = path
        self.device = device
        self.dtype = dtype
        self.iou_thres = iou_thres
        self.conf_thres = conf_thres
        self.imgsz = imgsz
        self.model = None

    def run(self) -> None:
        self.parent.start_loading_signal.emit('正在初始化模型')
        self.model = SingleDetect(self.parent, self.path, self.device, self.dtype, imgsz=self.imgsz)
        self.model.iou_thres = self.iou_thres
        self.model.conf_thres = self.conf_thres
        self.model.warmup()
        self.model.show_message = self.parent.output_print.emit  # 让detector可以输出信息
        self.parent.detector = self.model  # 预热后再替换
        self.parent.stop_loading_signal.emit('模型初始化成功')
        self.parent.show_detect_result.emit('')
